from brisque_dmos import *
import time
from sklearn.svm import SVR
from sklearn.model_selection import GridSearchCV
import joblib


start = time.perf_counter()
CreatSet()
creat_end = time.perf_counter()
creat_time = creat_end - start
print('The time of extract features of images: ')
print(creat_time)
print(creat_time / 982.0)

feat = np.load('feat_train.npy')
score = np.load('score_train.npy')
# search the value of C and γ from 1e-4 to 1e4 and take 9 candidate values
parameters = {"kernel": ("linear", 'rbf'), "C": np.logspace(-4, 4, 9), "gamma": np.logspace(-4, 4, 9)}
svr = GridSearchCV(SVR(), param_grid=parameters, cv=6)  # 6 cross validations
svr.fit(feat, score)
joblib.dump(svr, "dmos_model.m")
train_end = time.perf_counter()
train_time = train_end - creat_end
print('The time of train: ', train_time)
print('The parameters of the best model are: ')
print(svr.best_params_)
